package com.chief.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

public class OperatorStateTest {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<String> streamSource = env.socketTextStream("hadoop55", 7777);

        SingleOutputStreamOperator<Integer> map = streamSource.map(new MyOperatorMapFunction(33.33));

        map.print();

        env.execute("OperatorStateTest");
    }

    //为了故障恢复，还需要实现下列接口CheckpointedFunction
    public static class MyOperatorMapFunction implements MapFunction<String, Integer>, CheckpointedFunction {

        private Integer cnt = 0;

        private transient ListState<Integer> listState;

        private final Double level;

        public MyOperatorMapFunction(double v) {
            level = v;
        }

        @Override
        public Integer map(String value) throws Exception {
            String[] split = value.split(",");
            if (level.compareTo(Double.valueOf(split[2])) < 0) {
                cnt++;
            }
            return cnt;
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            listState.clear();
            listState.add(cnt);
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            listState = context.getOperatorStateStore().getListState(new ListStateDescriptor<>("list-state", Integer.class));
            for (Integer data : listState.get()) {
                cnt += data;
            }
        }
    }
}
